from ultralytics import YOLO
import argparse

def train_model(data_path, epochs=50, imgsz=640, batch=16, device='0'):
    # 加载YOLOv8-pose模型
    # 可选模型: yolov8n-pose.pt, yolov8s-pose.pt, yolov8m-pose.pt, yolov8l-pose.pt, yolov8x-pose.pt
    model = YOLO('yolov8s-pose.pt')  # 使用中等大小的模型
    
    # 开始训练
    results = model.train(
        data=data_path,        # 数据集配置文件路径
        epochs=epochs,         # 训练轮数
        imgsz=imgsz,           # 输入图像大小
        batch=batch,           # 批次大小
        device=device,         # 训练设备，'0'表示第一块GPU，'cpu'表示CPU
        task='pose',           # 任务类型：姿态估计
        mode='train',          # 模式：训练
        patience=10,           # 早停 patience，10轮没有提升就停止
        save=True,             # 保存模型
        save_period=10,        # 每10轮保存一次
        pretrained=True,       # 使用预训练权重
        optimizer='Adam',      # 优化器
        lr0=0.001,             # 初始学习率
        augment=True,          # 使用数据增强
        workers=4,             # 数据加载线程数
        project='card_corner_detection',  # 项目名称
        name='exp'             # 实验名称
    )
    
    return results

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='训练YOLOv8-pose卡片四角检测模型')
    parser.add_argument('--data', type=str, default='card_corner_dataset/data.yaml', help='数据集配置文件路径')
    parser.add_argument('--epochs', type=int, default=50, help='训练轮数')
    parser.add_argument('--imgsz', type=int, default=640, help='输入图像大小')
    parser.add_argument('--batch', type=int, default=16, help='批次大小')
    parser.add_argument('--device', type=str, default='0', help='训练设备')
    
    args = parser.parse_args()
    
    # 开始训练
    train_model(
        data_path=args.data,
        epochs=args.epochs,
        imgsz=args.imgsz,
        batch=args.batch,
        device=args.device
    )
    